import redis
from redisgraph import Graph
import torch
from tqdm import tqdm
import numpy as np

class DataLoaderToRedis:
    def __init__(self, redis_host, redis_port, graph_name, dataset_path, batch_size=1000):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.graph_name = graph_name
        self.dataset_path = dataset_path
        self.batch_size = batch_size
        self.graph = Graph(graph_name, redis.StrictRedis(host=redis_host, port=redis_port, db=0))
        self.dataset = torch.load(dataset_path)[0]
        self.features = self.dataset['x']
        self.edges = self.dataset['edge_index']
        # self.labels = self.dataset['y']
        self.labels = self.dataset['y'].argmax(dim = 1)
        
    def load_data_to_redis(self):
        """将数据加载到 RedisGraph 中，插入节点和边"""
        self._add_nodes()
        self._add_edges()

    def _add_nodes(self):
        """批量插入节点到 RedisGraph"""
        total_nodes = self.features.shape[0]
        
        for i in tqdm(range(0, total_nodes, self.batch_size), desc="Adding nodes"):
            batch_indices = range(i, min(i + self.batch_size, total_nodes))
            params_list = [
                {
                    'id': int(idx),
                    'feature': self.features[idx].tolist(),
                    'label': int(self.labels[idx].item())
                } 
                for idx in batch_indices
            ]
            
            query = """
            UNWIND $batch AS node
            MERGE (n:Node {id: node.id})
            SET n.feature = node.feature, n.label = node.label
            """
            self.graph.query(query, {'batch': params_list})
        
        print("============================节点添加完成============================")

    def _add_edges(self):
        """批量插入边"""
        edges_np = self.edges.t().numpy()
        total_edges = edges_np.shape[0]
        batch_size = 5000  # 增大批次大小

        for i in tqdm(range(0, total_edges, batch_size), desc="Adding edges"):
            batch = edges_np[i:i + batch_size]
            params = [{'src': int(src), 'dest': int(dst)} for src, dst in batch]
            
            # 优化查询：直接通过节点ID创建关系（假设节点已存在）
            query = """
            UNWIND $params AS edge
            CREATE (:Node {id: edge.src})-[:CONNECTED_TO]->(:Node {id: edge.dest})
            """
            try:
                self.graph.query(query, {'params': params})
            except redis.exceptions.ResponseError:
                # 若节点不存在，回退到MATCH-CREATE
                query_fallback = """
                UNWIND $params AS edge
                MATCH (a:Node {id: edge.src}), (b:Node {id: edge.dest})
                CREATE (a)-[:CONNECTED_TO]->(b)
                """
                self.graph.query(query_fallback, {'params': params})

class EmbeddingLoader:
    def __init__(self, redis_host, redis_port, graph_name, emb_file_path, batch_size=100):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.graph_name = graph_name
        self.emb_file_path = emb_file_path
        self.batch_size = batch_size
        self.graph = Graph(graph_name, redis.StrictRedis(host=redis_host, port=redis_port, db=0))

    def load_embeddings(self):
        """加载 .emb 文件并返回字典 {node_id: embedding_vector}"""
        embeddings = {}
        with open(self.emb_file_path, 'r') as f:
            header = f.readline().strip()
            num_nodes, dim = map(int, header.split())

            for line in tqdm(f, total=num_nodes, desc="Loading embeddings"):
                parts = line.strip().split()
                node_id = int(parts[0])
                vector = list(map(float, parts[1:]))

                assert len(vector) == dim, f"维度错误: 节点 {node_id} 的维度是 {len(vector)} (应为 {dim})"
                embeddings[node_id] = vector
        return embeddings

    def update_graph_embeddings(self, embeddings):
        """批量更新 RedisGraph 节点属性"""
        items = list(embeddings.items())
        
        for i in tqdm(range(0, len(items), self.batch_size), desc="Updating graph"):
            batch = items[i:i + self.batch_size]
            params = {
                'updates': [
                    {'node_id': node_id, 'embedding': vector}
                    for node_id, vector in batch
                ]
            }

            # query = """
            # UNWIND $updates AS update
            # MATCH (n:Node {id: update.node_id})
            # SET n.embedding = update.embedding
            # """
            query = """
            UNWIND $updates AS update
            MATCH (n) WHERE n.`id:ID` = update.node_id
            SET n.embedding = update.embedding
            """
            self.graph.query(query, params)

    def verify_embeddings(self, test_id):
        """验证更新的嵌入"""
        # query = """
        # MATCH (n:Node {id: $test_id})
        # RETURN n.id AS node_id, size(n.embedding) AS dim
        # """
        query = """
        MATCH (n) WHERE n.`id:ID` = $test_id
        RETURN n.`id:ID` AS node_id, n.embedding AS embedding
        """
        params = {'test_id': test_id}
        res = self.graph.query(query, params)
        print(f"验证节点 {test_id}:", res.result_set)